import numpy as np
import math
def edgeDetectKernels(kernel_type, sigma = 1):
    """
    2D Gause Kernel.
    """
    SUM = 0.0
    xtmp = [[1,0,-1],
            [2,0,-2],
            [1,0,-1]]
    ytmp = [[1,2,1],
            [0,0,0],
            [-1,-2,-1]]
    xKernel = np.array(xtmp)
    yKernel = np.array(ytmp)
    center = 1
    if kernel_type == "gause derivative":
        for i in range(3):
            for j in range(3):
                x = i - center
                y = j - center
                xKernel[i][j] = (-x / sigma**2) * math.exp( - (x**2 + y**2) / (2 * sigma**2))
                yKernel[i][j] = (-y / sigma**2) * math.exp( - (x**2 + y**2) / (2 * sigma**2))
    elif kernel_type == "sobel" or kernel_type == "Sobel":
        return xKernel, yKernel
    else:
        raise Exception("kernelTypeError, please enter 'gause derivative' or 'sobel'.")
    return xKernel, yKernel

class Conv2D():
    """
    A very simple implement of convolutional layer.
    """

    def __init__(self, filters, kernel_size, sigma, boundary = "zero", padding = "same"):
        xKernel, yKernel = edgeDetectKernels("Sobel", sigma)
        self.filters = [xKernel, yKernel]
        self.kernel_size = kernel_size
        self.sigma = sigma
        self.padding = padding
        self.outputs = []
        self.boundary = boundary
    
    def subOfTwoKernels(self, inputs, kernel1, kernel2):
        self.filters = [kernel1, kernel2]
        outputs = self.call(inputs)
        return outputs[0] - outputs[1]

    def boundaryChange(self, inputs):
        image = np.array(inputs)
        h = image.shape[0]
        w = image.shape[1]
        
        imageExtend = np.zeros(shape = (3*h,3*w))
        for i in range(h, 2*h):
            for j in range(w, 2*w):
                imageExtend[i][j] = image[i-h][j-w]
        if self.boundary == "wrap":
            for i in range(3*h):
                for j in range(3*w):
                    imageExtend[i][j] = image[i % h][j % w]
        elif self.boundary == "copy":
            for i in range(h):
                for j in range(w,2*w):
                    imageExtend[i][j] = imageExtend[h][j]
            for i in range(2*h, 3*h):
                for j in range(w,2*w):
                    imageExtend[i][j] = imageExtend[2*h-1][j]
            for i in range(h,2*h):
                for j in range(w):
                    imageExtend[i][j] = imageExtend[i][w]
            for i in range(h,2*h):
                for j in range(2*w,3*w):
                    imageExtend[i][j] = imageExtend[i][2*w-1]
            for i in range(h):
                for j in range(w):
                    imageExtend[i][j] = imageExtend[h][w]
            for i in range(h):
                for j in range(2*w, 3*w):
                    imageExtend[i][j] = imageExtend[h][2*w-1]
            for i in range(2*h,3*h):
                for j in range(w):
                    imageExtend[i][j] = imageExtend[2*h-1][w]
            for i in range(2*h,3*h):
                for j in range(2*w,3*w):
                    imageExtend[i][j] = imageExtend[2*h-1][2*w-1]
        elif self.boundary == "reflect":
            imageUpDown = np.flip(image, axis=0)
            imageLeftRight = np.flip(image, axis=1)
            imageDiag = np.rot90(image, 2)
            for i in range(h):
                for j in range(w):
                    imageExtend[i][j] = imageDiag[i][j]
            for i in range(2*h, 3*h):
                for j in range(2*w, 3*w):
                    imageExtend[i][j] = imageDiag[i-2*h][j-2*w]
            for i in range(h, 2*h):
                for j in range(w):
                    imageExtend[i][j] = imageLeftRight[i-h][j]
            for i in range(h,2*h):
                for j in range(2*w, 3*w):
                    imageExtend[i][j] = imageLeftRight[i-h][j-2*w]
            for i in range(h):
                for j in range(w, 2*w):
                    imageExtend[i][j] = imageUpDown[i][j-w]
            for i in range(2*h, 3*h):
                for j in range(w, 2*w):
                    imageExtend[i][j] = imageUpDown[i-2*h][j-w]
            for i in range(h):
                for j in range(2*w, 3*w):
                    imageExtend[i][j] = imageDiag[i][j-2*w]
            for i in range(2*h, 3*h):
                for j in range(w):
                    imageExtend[i][j] = imageDiag[i-2*h][j]
            
        outputImage = Image.fromarray(imageExtend.astype("uint")).convert("L")
        name = "./GK/rawImageWith" + self.boundary + ".jpg"
        outputImage.save(name)
        return imageExtend
                    

    def call(self, inputs):
        inputs = self.boundaryChange(inputs)
        image = np.array(inputs)
        h = len(image) // 3
        w = len(image[0]) // 3
        if self.padding == "SAME" or self.padding == "same":
            for f in range(len(self.filters)):
                output = np.zeros(shape = (h, w))
                for i in range(h):
                    for j in range(w):
                        output[i][j] = self.multiply_element_wide(i, j, h, w, image, self.filters[f])
                self.outputs.append(output)
        elif self.padding == "VALID" or self.padding == "valid":
            for f in range(len(self.filters)):
                k = self.kernel_size // 2
                output = np.zeros(shape = (h - 2*k, w - 2*k))
                for i in range(k, h - k):
                    for j in range(k, w - k):
                        output[i - k][j - k] = self.multiply_element_wide(i, j, h, w, image, self.filters[f])
                self.outputs.append(output)
        elif self.padding == "FULL" or self.padding == "full":
            for f in range(len(self.filters)):
                k = self.kernel_size // 2
                output = np.zeros(shape = (h + 2*k, w + 2*k))
                for i in range(- k, h + k):
                    for j in range(- k, w + k):
                        output[i + k][j + k] = self.multiply_element_wide(i, j, h, w, image, self.filters[f])
                self.outputs.append(output)
        else:
            raise Exception("PaddingError, please enter valid or same.")
        return self.outputs

    def multiply_element_wide(self, i, j, h, w, image, kernel):
        SUM = 0.0
        k = self.kernel_size // 2
        center = k
        for m in range(-k, k+1):
            for n in range(-k, k+1):
                # if i + m >= 0 and j + n >= 0 and i + m < h and j + n < w:
                tmp = kernel[center + m][center + n] * image[h + i + m][w + j + n]
                SUM += tmp
        return SUM
    
    def __call__(self, inputs, **kwds):
        return self.call(inputs)
    def getMagnitudeAndOrientation(self):
        xGrad, yGrad = self.xGrad, self.yGrad
        h = len(xGrad)
        w = len(xGrad[0])
        self.h = h
        self.w = w
        magnitude = np.zeros(shape = (h, w))
        orientation = np.zeros(shape = (h, w, 2))
        for i in range(h):
            for j in range(w):
                magnitude[i][j] = sqrt(xGrad[i][j]**2 + yGrad[i][j]**2) + 0.000000001
                orientation[i][j][0] = yGrad[i][j] / magnitude[i][j]
                orientation[i][j][1] = xGrad[i][j] / magnitude[i][j]
        self.magnitude = magnitude
        self.orientation = orientation
        MagnitudeImage = Image.fromarray(self.magnitude.astype("uint")).convert("L")
        MagnitudeImage.save("./EdgeDetect/magnitude.jpg")
        return magnitude, orientation